// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================

/*!
\file
\brief Implementation of Strata or Layer (SOL) Utility Functions 
*/

#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "boolfile.h"
#include "copyfile.h"
#include "phefile.h"
#include "pmaths.h"
#include "sample.h"
#include "solfile.h"
#include "suffixes.h"
#include "utils.h"

#define ID_COL_INDEX 0
#define MAX_LINE_LENGTH 200000

int _get_levels(char **src, int nSrc, char **dst, int nDst, int *nLevels, BOOL allowEmptyStrings, BOOL *items_to_ignore) {
	int i, j, levelCount = 0;
	
	if (src == NULL || nSrc < 0 || dst == NULL || nDst < 0 || nLevels == NULL || items_to_ignore == NULL) return 0;

	set_all_lowercase(src, nSrc);
	zero_2d_char_array(dst, nDst);
	*nLevels = 0;

	// copy 1st element from the src array into dst as need something
	// to  start with when getting the levels from the array.
	for (i = 0; i < nSrc; i++) {
		if (strlen(src[i]) > 0 && strcmp(src[i], MISSING_LEVEL) != 0 && !items_to_ignore[i]) {
			strcpy(dst[0], src[i]);
			levelCount++;
			break;
		}
	}
	if (levelCount == 0) return 0; // no initial level so 'nuts' to carrying on.

	// Compare all values to get the levels, if the maximum values reached assume
	// column is a continuous data type.
	for (i = 0; i < nSrc; i++) {
		BOOL isFound = FALSE;

		if (levelCount == nDst) break; // Hit max. levels.
		if (!allowEmptyStrings) {
			if (strlen(src[i]) == 0) continue;
		}
		if(strcmp(src[i], MISSING_LEVEL) == 0) continue;

		for (j = 0; j < levelCount; j++) {
			if(strcmp(src[i], dst[j]) == 0) {
				isFound = TRUE;
				break;
			}
		}
		if (!isFound) {
			strcpy(dst[levelCount], src[i]);
			levelCount++;
		}
	}
	*nLevels = levelCount;

	// If level count hit's max, this means that the column value ids likely to be
	// a continous field and so doesn't have levels.
	if (levelCount == nDst) {
		*nLevels = 0;
		return 0;
	} else {
		qsort(dst,levelCount, sizeof(char*),cstring_cmp);
		return 1;
	}
}

#ifndef MAX_LINE_WIDTH
#define MAX_LINE_WIDTH 131072
#endif

int solfile_create(struct selected_options *ops, BOOL skip_missing) {
	BOOL jobDone = FALSE, *persons_to_del = NULL, hasHeader = FALSE;
	int i, j, k, nSamples = 0, nActualLevels = 0;
	char localFilepath[FILE_PATH_LENGTH], remoteFilepath[FILE_PATH_LENGTH], **levels = NULL, **values = NULL, **ids = NULL, **lines = NULL;
	int **membership_data = NULL, *membership_counts = NULL, filtered_dataset_row_index = 0, nLines = 0;
	size_t nPersonsToKeep = 0; 
	
	if (ops == NULL || ops->mpheno < 0 || skip_missing != TRUE && skip_missing != FALSE) return 0;
	else if (strlen(ops->szPheFileName) == 0 || strlen(ops->ldir) == 0 || strlen(ops->rdir) == 0 || strlen(ops->szRunId) == 0 || strlen(ops->ldir) == 0) return 0;
	else if (ops->nperm <= 0) return 0;
	
	// Convert the PHE-PLINK index to actual token index.
	ops->mpheno = ops->mpheno - 1 + 2;

	// Load the entire COVAR file into memory, wasteful but this code may be converted to multi-column variants 
	// in case the code is in situ for the refactoring.
	strcpy(remoteFilepath, ops->szPheFileName);
	sprintf(localFilepath, "%s%s%s", ops->ldir, ops->szRunId, SUFFIX_PHE_FILE);
	
	// Load the PTD file of people still in the analysis data-set after filtering.
	persons_to_del = boolfile_read(ops, SUFFIX_PERSON_TO_DEL, &nSamples);
	if (nSamples == 0 || persons_to_del == NULL) return 0;
	for (i = 0; i < nSamples; i++) if (!persons_to_del[i]) nPersonsToKeep++;
	if (nPersonsToKeep == 0) { // nobody left in the data-set so exit early.
		if (persons_to_del != NULL) free(persons_to_del);
		return 0;
	}
	
	// Load the column from the input PHE file containing all of the reference information.
	if (copyFile(remoteFilepath, localFilepath)) {
		// Find index of a named column.
		if (strlen(ops->szSolColumnName) > 0) {
			FILE *in = fopen(localFilepath, "r");
			if (in != NULL) {
				char line[MAX_LINE_LENGTH], token[TOKEN_LENGTH];
				if(fgets(line,MAX_LINE_LENGTH, in)) {
					int nTokens = get_number_of_tokens(line, ' '), i;
					for (i = 0; i < nTokens; i++) {
						get_token(line, token, i, ' ');
						if (strcmp(token, ops->szSolColumnName) == 0) {
							hasHeader = TRUE; 
							ops->mpheno = i;
							break;
						}
					}
				}
				fclose(in);
				if (!hasHeader) {
					if (persons_to_del != NULL) free(persons_to_del);
					deleteFile(localFilepath);
					return 0;
				} 
			}
		}
		
		values = create_2d_char_array(nSamples, MAX_LEVEL_WIDTH);
		ids = create_2d_char_array(nSamples, 20);
		lines = load_file(localFilepath, &nLines);
		if (lines != NULL && nLines > 0) {
			int pheIndex = 0;

			// Skip the header if present.
			if (hasHeader) j = 1;
			else j = 0;
			for (i = j, pheIndex = 0; i < nLines; i++, pheIndex++) {
				get_token(lines[i], ids[pheIndex], ID_COL_INDEX, ' ');
				get_token(lines[i], values[pheIndex], ops->mpheno, ' ');
			}
			jobDone = TRUE;
		}
		deleteFile(localFilepath);
		if (jobDone) {
			jobDone = FALSE;
			levels = create_2d_char_array(MAX_LEVELS, MAX_LEVEL_WIDTH);
			if (_get_levels(values, nSamples, levels, MAX_LEVELS, &nActualLevels, FALSE, persons_to_del)) {
				if (nActualLevels >= 2) {
					FILE *out = NULL;

					sprintf(localFilepath, "%s%s%s", ops->ldir, ops->szRunId, SUFFIX_SOL_FILE);
					sprintf(remoteFilepath, "%s%s%s", ops->rdir, ops->szRunId, SUFFIX_SOL_FILE);

					out = fopen(localFilepath, "w");
					membership_counts = create_int_array(nActualLevels, 0);
					membership_data = create_2d_int_array(nActualLevels, nPersonsToKeep, -1);

					if (out != NULL && membership_counts != NULL && membership_data != NULL) {
						for (i = 0; i < nSamples; i++) {
							BOOL is_missing_level = FALSE;
							int level_index = 0;

							if (persons_to_del[i]) continue;
							if (strcmp(values[i], MISSING_LEVEL) == 0) is_missing_level = TRUE;

							if (!is_missing_level) {
								for (j = 0; j < nActualLevels; j++) {
									if (strcmp(values[i], levels[j]) == 0) {
										membership_data[j][membership_counts[j]++] = filtered_dataset_row_index; 
										level_index = j;
										break;
									}
								}
							} 

							if (is_missing_level) {
								if (!skip_missing) fprintf(out, "%s 1 %s\n", ids[i], MISSING_LEVEL);
							} else fprintf(out, "%s 1 %i\n", ids[i], level_index);

							filtered_dataset_row_index++; // A person included in the filtered data-set, 'i' in this loop is the unfiltered array indices.
						}
						fclose(out);
						if (copyAndDeleteFile(localFilepath, remoteFilepath)) jobDone = TRUE;
					}
				}
			}
		}
	}

	// Write out the set permutation file.
	if (jobDone == TRUE) {
		int *full_perm_row = NULL, *set_perm_row = NULL, nMax_level_member = 0;
	
		jobDone = FALSE;

		// Get the max level size in terms of number of members.
		for (i = 0; i < nActualLevels; i++) nMax_level_member = max(nMax_level_member, membership_counts[i]);
		
		full_perm_row = create_int_array(nPersonsToKeep, 0);
		set_perm_row = create_int_array(nMax_level_member, 0);

		if (full_perm_row != NULL && set_perm_row != NULL) {
			int permsRowsWrote = 0;
			FILE *out = NULL;

			sprintf(localFilepath, "%s%s%s", ops->ldir, ops->szRunId, SUFFIX_PERM_FILE);
			sprintf(remoteFilepath, "%s%s%s", ops->rdir, ops->szRunId, SUFFIX_PERM_FILE);

			out = fopen(localFilepath, "wb");
			if (out != NULL) {
				if (fwrite(&(ops->nperm), sizeof(int), 1, out) == 1) {
					if (fwrite(&nPersonsToKeep, sizeof(int), 1, out) == 1) {
						for (i = 0; i < ops->nperm; i++) {
							fill_indices(full_perm_row, nPersonsToKeep);
							for (j = 0; j < nActualLevels; j++) {
								permute(set_perm_row, membership_counts[j]);
								for (k = 0; k < membership_counts[j]; k++) full_perm_row[membership_data[j][k]] = membership_data[j][set_perm_row[k]];
							}
							if (fwrite(full_perm_row, sizeof(int), nPersonsToKeep, out) == nPersonsToKeep) permsRowsWrote++;
							else break;
						}
					}
				}
				fclose(out);
				if(copyAndDeleteFile(localFilepath, remoteFilepath)) {
					if (permsRowsWrote == ops->nperm) jobDone = TRUE;
				}
			}
		}

		if (full_perm_row != NULL) free(full_perm_row);
		if (set_perm_row != NULL) free(set_perm_row);
	}

	// Free up stuff
	if (persons_to_del != NULL) free(persons_to_del);
	if (levels != NULL) free_2d_char_array(levels, MAX_LEVELS);
	if (ids != NULL) free_2d_char_array(ids, nSamples);
	if (values != NULL) free_2d_char_array(values, nSamples);
	if (membership_data != NULL) free_2d_int_array(membership_data, nActualLevels);
	if (membership_counts != NULL) free(membership_counts);
	if (lines != NULL) free_2d_char_array(lines, nLines);

	if (jobDone) return 1;
	else return 0;
}

int solfile_load(struct selected_options *ops, struct sample *samples, int nSamples) {
	BOOL jobDone = FALSE, *missing = NULL;
	double *values = NULL;
	char tmp_phe_filename[FILE_PATH_LENGTH]; // Preserve the PHE filename
	int MISSING = atoi(MISSING_LEVEL);

	*tmp_phe_filename = '\0';

	if (ops == NULL || samples == NULL || nSamples <= 0) return 0;
	else if (strlen(ops->szSolFileName) == 0) return 0;

	ops->nk = 0;
	strcpy(tmp_phe_filename, ops->szPheFileName);
	strcpy(ops->szPheFileName, ops->szSolFileName);
	
	missing = create_bool_array(nSamples, FALSE);
	if (missing != NULL) {
		values = phefile_load(ops, nSamples, missing);
		if (values != NULL) {
			int i;
			for (i = 0; i < nSamples; i++) {
				int value = (int) values[i];
				if (value != MISSING) samples[i].sol = value;
			}
			jobDone = TRUE;

			// Get the number of levels.
			for (i = 0; i < nSamples; i++) ops->nk = max(ops->nk, samples[i].sol);
			ops->nk++; // Sol is an index field, not a count.
		}
	}
	strcpy(ops->szPheFileName, tmp_phe_filename);
	if (missing != NULL) free(missing);
	if (values != NULL) free(values);
	
	if (jobDone) return 1; 
	else return 0;
}
